import open3d as o3d
import numpy as np
from scipy import optimize

def showPCDRGB(l_p,xyz,rgb,size_frame=1):
    # vis = o3d.visualization.Visualizer()
    # vis.create_window()
    point_cloud_o3d = o3d.geometry.PointCloud()
    point_cloud_o3d.points = o3d.utility.Vector3dVector(xyz[:,:3])
    point_cloud_o3d.colors = o3d.utility.Vector3dVector(rgb/255)
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size_frame)
    # vis.add_geometry(point_cloud_o3d)
    # vis.add_geometry(frame)
    pcd = o3d.io.read_point_cloud(l_p)

    vis_add = o3d.geometry.PointCloud()
    vis_add.points = o3d.utility.Vector3dVector([[0,0,1]])
    vis_add.colors = o3d.utility.Vector3dVector([[0,0,1]])

    o3d.visualization.draw_geometries([point_cloud_o3d,frame,pcd])
    # vis.run()
    # vis.destroy_window()

def r(x, y, xc, yc):
    return np.sqrt((x-xc)**2 + (y-yc)**2)

def f(c, x, y):
    Ri = r(x, y, *c)
    return np.square(Ri - Ri.mean())

def least_squares_circle(coords):
    """
    Circle fit using least-squares solver.
    Inputs:

        - coords, list or numpy array with len>2 of the form:
        [
    [x_coord, y_coord],
    ...,
    [x_coord, y_coord]
    ]
        or numpy array of shape (n, 2)

    Outputs:

        - xc : x-coordinate of solution center (float)
        - yc : y-coordinate of solution center (float)
        - R : Radius of solution (float)
        - residu : MSE of solution against training data (float)
    """

    x, y = None, None
    if isinstance(coords, np.ndarray):
        x = coords[:, 0]
        y = coords[:, 1]
    elif isinstance(coords, list):
        x = np.array([point[0] for point in coords])
        y = np.array([point[1] for point in coords])
    else:
        raise Exception("Parameter 'coords' is an unsupported type: " + str(type(coords)))

    # coordinates of the barycenter
    x_m = np.mean(x)
    y_m = np.mean(y)
    center_estimate = x_m, y_m
    center, _ = optimize.leastsq(f, center_estimate, args=(x, y))
    xc, yc = center
    Ri       = r(x, y, *center)
    R        = Ri.mean()
    residu   = np.sum((Ri - R)**2)
    return xc, yc, R, residu

def get_center_xyz(cylinder_pl,pcd):
    # cylinder_pl = [2.5,3.5]
    z_range=[0.2,2]

    x_range = np.array([cylinder_pl[0]-0.3,cylinder_pl[0]+0.3])
    y_range = np.array([cylinder_pl[1]-0.3,cylinder_pl[1]+0.3])
    

    crop_box = o3d.geometry.AxisAlignedBoundingBox([x_range[0], y_range[0], z_range[0]], [x_range[1],y_range[1], z_range[1]])
    pcd_cropped = pcd.crop(crop_box)

    o3d.visualization.draw_geometries([pcd_cropped])

    # cylinder
    xyz = np.asarray(pcd_cropped.points)
    a, b, r, residual = least_squares_circle(xyz[:,:2])
    print('Center: [{}, {}, {}], Radius: {}'.format(a, b, 0, r))
    cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=r, height=1)

    cylinder.translate([a,b,0])

    o3d.visualization.draw_geometries([pcd_cropped,cylinder])
    print("xy: ", [a,b])
    return a,b

def pick_points(pcd,size_frame=0.3):
    print("")
    print(
        "1) Please pick at least three correspondences using [shift + left click]"
    )
    print("   Press [shift + right click] to undo point picking")
    print("2) After picking points, press 'Q' to close the window")
    vis = o3d.visualization.VisualizerWithEditing()
    vis.create_window(window_name="pick pts",visible=True)
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size_frame)
    vis.add_geometry(pcd)
    vis.add_geometry(frame)
    # o3d.visualization.draw_geometries([point_cloud_o3d,frame])
    
    vis.run()  # user picks points
    vis.destroy_window()
    print("")
    return vis.get_picked_points()

def show_plane(a,b,c,d,crop_pcd):

    plane = o3d.geometry.TriangleMesh()

    vertices = [
        [100.0, 100.0, (-d - a - b)],
        [-100.0, 100.0, (-d + a + b)],
        [-100.0, -100.0, (-d + a - b)],
        [100.0, -100.0, (-d - a + b)]
    ]
    plane.vertices = o3d.utility.Vector3dVector(vertices)

    triangles = [
        [0, 1, 2],
        [0, 2, 3]
    ]
    plane.triangles = o3d.utility.Vector3iVector(triangles)
    plane.paint_uniform_color([0, 1, 0])
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)

    vis = o3d.visualization.Visualizer()
    vis.create_window()
    vis.add_geometry(plane)
    vis.add_geometry(frame)
    vis.add_geometry(crop_pcd)

    vis.run()
    vis.destroy_window()


def get_plane(pcd):
    bbox=o3d.geometry.AxisAlignedBoundingBox([-50, -20, -0.5], [100,20, 0.5])
    clipped_pcd = pcd.crop(bbox)

    plane_model, inliers = clipped_pcd.segment_plane(distance_threshold=0.02,
                                                    ransac_n=3,
                                                    num_iterations=1000)

    [a, b, c, d] = plane_model
    show_plane(a,b,c,d,clipped_pcd)
    print("plane: {}x + {}y + {}z + {} = 0".format(a, b, c, d))
    return a,b,c,d

def calculate_z_given_plane_equation(a, b, c, d, x, y):
    z = (-d - a*x - b*y) / c
    return z

def get_lidar_pts(lidar_p,pts_num):
    pcd = o3d.io.read_point_cloud(lidar_p)

    # ground fit , get z of p-cylinder in plane
    a,b,c,d = get_plane(pcd)

    # pts_num = 4 cylinder center
    lidar_pts = []
    for i in range(pts_num):
        pts_ind = pick_points(pcd)
        circle_x,circle_y,_ = pcd.points[pts_ind[0]]
        x,y = get_center_xyz([circle_x,circle_y],pcd)

        z = calculate_z_given_plane_equation(a,b,c,d,x,y)
        print("x,y,z \n",x,y,z)
        lidar_pts.append([x,y,z])

    return np.array(lidar_pts,dtype=np.float32)

if __name__=="__main__":
    lidar_p = "/home/westwell/welldriver/log/data_calib/qtruck_psa001/zjh_0307_snap/snap_rl/qtruck_psa001_rl/lidar_left_ego/1709787433725392000.pcd"
    pts_num = 2
    pts = get_lidar_pts(lidar_p,pts_num)
    print(pts)
    # [[2.4653146  3.6818726  0.        ]
    # [3.0435183  4.168749   0.        ]
    # [0.04880002 4.770845   0.        ]
    # [2.8511071  6.6054993  0.        ]]
    




